from config.settings import ZOOMEYE_SAVE_PATH, ZOOMEYE_AUTH
import json
import os
from zoomeye.sdk import ZoomEye
from export.export import export_to_excel
from lib import join

# default config
ZOOMEYE_USERNAME = ZOOMEYE_AUTH[0][0]
ZOOMEYE_PASSWORD = ZOOMEYE_AUTH[0][1]

class Counter(dict):
    """TODO: 修改嵌套逻辑使得 dict[one][two] 中one和two的key不存在时，返回空 """

    def __missing__(self, key):
        return ''


class ZoomeyeScan:
    """IP输出信息是一个端口对应一个对象
    """
    name = 'ZoomEye'

    def __init__(self):
        self.register()
        self.ip = ''

    def register(self):
        self.zm = ZoomEye(username=ZOOMEYE_USERNAME, password=ZOOMEYE_PASSWORD)
        info = self.zm.login()

    def ip_infos(self, ip):
        self.ip = ip
        if self.cache(ip):
            return self.format(ip)
        try:
            self.result = self.zm.dork_search(f'ip:{ip}', page=0, resource='host')
            if not self.result:
                self.result = {}
        except:
            self.result = {}

        self.save(ip)
        self.format(ip)
        return self.result

    def net_infos(self, network):
        if self.cache(network):
            return self.format(network)
        self.result = self.zm.dork_search(f'cidr:{network}', resource='host')
        self.save(network)
        return self.result

    def save(self, ip: str):
        if '/' in ip:
            ip = ip.replace('/', '_')
        with open(f'{ZOOMEYE_SAVE_PATH}/{ip}.json', 'w+') as f:
            json.dump(self.result, f)

    def cache(self, ip: str):
        if '/' in ip:
            ip = ip.replace('/', '_')
        if os.path.exists(f'{ZOOMEYE_SAVE_PATH}/{ip}.json'):
            with open(f'{ZOOMEYE_SAVE_PATH}/{ip}.json', 'r') as f:
                self.result = json.load(f)
                return True
        return False

    def format(self, ip) -> dict:
        results = {ip: []}
        for x in self.result:
            results[ip].append(x['portinfo'])
        return results

    @property
    def current(self):

        return {
            'application': '',
            'probe': '',
            'transport': '',
            'openport': '',
            'service': '',
            'title': '',
            'extrainfo': '',
            'device': '',
            'os': '',
            'version': '',
            'app': '',
            'banner': '',
            'icon_hash': ''

        }

    def parse(self, raw):
        current = self.current
        raw = Counter(raw)
        if 'protocol' in raw:
            current['application'] = raw['protocol']['application']
            current['probe'] = raw['protocol']['probe']
            current['transport'] = raw['protocol']['transport']
        current['ip'] = raw['ip']
        current['city'] = raw['geoinfo']['city']['names']['zh-CN']
        current['idc'] = raw['geoinfo']['idc']
        current['isp'] = raw['geoinfo']['isp']
        current['asn'] = raw['geoinfo']['asn']
        current['subdivisions'] = raw['geoinfo']['subdivisions']['names']['zh-CN']
        current['location'] = f'{raw["geoinfo"]["location"]["lat"]}, {raw["geoinfo"]["location"]["lon"]}'
        if 'portinfo' in raw:
            current['openport'] = raw['portinfo']['port']
            current['service'] = raw['portinfo']['service']
            current['title'] = ','.join(join(raw['portinfo']['title']))
            current['extrainfo'] = raw['portinfo']['extrainfo']
            current['device'] = raw['portinfo']['device']
            current['os'] = raw['portinfo']['os']
            current['version'] = raw['portinfo']['version']
            current['app'] = raw['portinfo']['app']
            current['banner'] = raw['portinfo']['banner']
        current['honeypot'] = raw['honeypot']
        if 'ico_mmh3' in raw:
            current['ico_mmh3'] = raw['ico']['mmh3']
        current['source_file'] = f'{ZOOMEYE_SAVE_PATH}/{self.ip}'
        current['engine'] = 'zoomeye'
        return current

    def export(self, save_path=None):
        export_results = []
        for x in self.result:
            current = self.parse(x)
            export_results.append(current)
        export_to_excel(export_results, sheet_name='IP信息', xlsx_name=f'{self.ip}', save_path=save_path)


class Zoomeye:
    def __init__(self, ip_list: list, query='zoomeye_search'):
        self.ip_list = ip_list
        self.scan = ZoomeyeScan()
        self.results = []
        self.query = query

    def run(self, save_path):
        for ip in self.ip_list:
            self.scan.ip_infos(ip)
            for x in self.scan.result:
                current = self.scan.parse(x)
                self.results.append(current)
        self.export(save_path)

    def export(self, save_path=None):
        export_to_excel(self.results, sheet_name='IP信息', xlsx_name=f'{self.query}', save_path=save_path)


